# -*- coding: utf-8 -*-
"""
@Time    : 2024/9/5 11:06 
@Author  : ZhangShenao 
@File    : live_agent.py 
@Desc    : 直播带货营销智能体
"""

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableConfig
from langgraph.graph import MessagesState, StateGraph
from langgraph.graph.state import CompiledStateGraph
from langgraph.prebuilt import ToolNode, tools_condition

from base import BaseAgentState, get_llm
from google_serper_tool import create_google_serper_tool


# 定义直播带货智能体状态
class LiveAgentState(MessagesState, BaseAgentState):
    """直播带货Agent状态"""
    pass


# 定义直播带货智能体节点
def live_agent_node(state: LiveAgentState, config: RunnableConfig) -> LiveAgentState:
    """直播带货智能体节点"""

    # 创建工具列表
    tools = [create_google_serper_tool()]

    # 获取LLM,并绑定工具列表
    llm = get_llm().bind_tools(tools=tools)

    # 构造Prompt
    prompt = ChatPromptTemplate.from_messages([
        ("system",
         "你是一位拥有10年经验的直播文案专家，请根据用户提供的产品整理一篇直播带货脚本文案，如果在你的知识库内找不到关于该产品的信息，可以使用搜索工具。"),
        ("human", "{query}"),
        ("placeholder", "{chat_history}"),  # 需要在Prompt中传入聊天历史记录,便于LLM从历史的ToolMessage中获取工具调用信息
    ])

    # 构造Chain
    chain = prompt | llm

    # 调用Chain,获取结果
    reply = chain.invoke({"query": state["query"], "chat_history": state["messages"]})

    # 返回更新状态
    return {
        "live_content": reply.content,
        "messages": reply,
    }


# 定义创建子图的函数
def create_live_graph() -> CompiledStateGraph:
    """创建直播带货子图"""

    # 创建图架构
    graph_builder = StateGraph(LiveAgentState)

    # 添加直播Agent节点
    graph_builder.add_node(node="live_agent", action=live_agent_node)
    # 添加工具调用节点,并指定工具列表
    graph_builder.add_node(node="tools", action=ToolNode(tools=[create_google_serper_tool()]))

    # 设置入口节点
    graph_builder.set_entry_point(key="live_agent")
    # 添加条件边,采用LangGraph内置的tools_condition条件函数
    # 根据LLM的生成结果中是否包含工具调用信息,决定是路由到tools节点,还是直接路由到END节点
    graph_builder.add_conditional_edges(source="live_agent", path=tools_condition)
    graph_builder.add_edge(start_key="tools", end_key="live_agent")

    # 将图编译成Runnable可执行组件,并返回
    return graph_builder.compile()
